# -*- coding: utf-8 -*-
import math
import tf.transformations as tft
import sys
import rospy
import moveit_commander
import geometry_msgs.msg

def euler_to_quaternion(roll, pitch, yaw):
    # Convert angles from degrees to radians
    roll_rad = math.radians(roll)
    pitch_rad = math.radians(pitch)
    yaw_rad = math.radians(yaw)

    # Compute quaternion
    q = tft.quaternion_from_euler(roll_rad, pitch_rad, yaw_rad, axes='sxyz')
    return q

def compute_ik(x, y, z, r_x, r_y, r_z, r_w):
    # 初始化 ROS 节点和 MoveIt
    rospy.init_node('compute_ik', anonymous=True)
    moveit_commander.roscpp_initialize(sys.argv)
    arm = moveit_commander.MoveGroupCommander("manipulator")

    # 设置目标位姿
    target_pose = geometry_msgs.msg.Pose()
    target_pose.position.x = x
    target_pose.position.y = y
    target_pose.position.z = z
    target_pose.orientation.x = r_x
    target_pose.orientation.y = r_y
    target_pose.orientation.z = r_z
    target_pose.orientation.w = r_w

    # 设置目标位姿
    arm.set_pose_target(target_pose)
    arm.set_planning_time(10.0)

    # 规划
    success = arm.plan()
    if success:
        # 获取关节值
        joint_values = arm.get_current_joint_values()
        rospy.loginfo("Joint values for look position: %s", joint_values)
    else:
        rospy.logwarn("Failed to compute IK")

    moveit_commander.roscpp_shutdown()

if __name__ == '__main__':
    # 获取用户输入（兼容 Python 2 和 3）
    try:
        if sys.version_info[0] < 3:
            roll = float(raw_input('Rx: '))
            pitch = float(raw_input('Ry: '))
            yaw = float(raw_input('Rz: '))
        else:
            roll = float(input('Rx: '))
            pitch = float(input('Ry: '))
            yaw = float(input('Rz: '))
    except ValueError:
        print("Error: Please enter valid numbers for Rx, Ry, and Rz")
        sys.exit(1)

    # Convert to quaternion
    quaternion = euler_to_quaternion(roll, pitch, yaw)
    print("Quaternion:", quaternion)
    print("Quaternion type:", type(quaternion))

    # 询问用户是否计算 IK
    compute = raw_input("Compute IK? (y/n): ") if sys.version_info[0] < 3 else input("Compute IK? (y/n): ")
    if compute.lower() == 'y':
        try:
            x = float(raw_input('x: ') if sys.version_info[0] < 3 else input('x: '))
            y = float(raw_input('y: ') if sys.version_info[0] < 3 else input('y: '))
            z = float(raw_input('z: ') if sys.version_info[0] < 3 else input('z: '))
            compute_ik(x, y, z, quaternion[0], quaternion[1], quaternion[2], quaternion[3])
        except ValueError:
            print("Error: Please enter valid numbers for x, y, and z")
            sys.exit(1)
        except rospy.ROSInterruptException:
            pass